#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
离线海图下载修复工具

这个脚本用于修复失败的下载任务，尝试使用不同的瓦片源再次下载缺失的瓦片。
"""

import os
import sys
import argparse
import logging
import json
from tqdm import tqdm

# 添加项目根目录到系统路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = current_dir
sys.path.append(project_root)

# 导入自定义模块
from backend.downloader import TileDownloader, TILE_SOURCES
from backend.converter import MBTilesConverter

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('RepairTool')

def list_task_dirs():
    """列出所有任务目录"""
    data_dir = os.path.join(project_root, "data", "tiles")
    if not os.path.exists(data_dir):
        logger.error(f"数据目录不存在: {data_dir}")
        return []
    
    task_dirs = []
    for item in os.listdir(data_dir):
        if item.startswith("task_") and os.path.isdir(os.path.join(data_dir, item)):
            task_dirs.append(item)
    
    return sorted(task_dirs, key=lambda x: int(x.split('_')[1]))

def scan_task_dir(task_dir_name):
    """扫描任务目录中的瓦片文件"""
    task_dir = os.path.join(project_root, "data", "tiles", task_dir_name)
    tiles = []
    
    # 遍历目录，统计瓦片文件
    for zoom_dir in os.listdir(task_dir):
        zoom_path = os.path.join(task_dir, zoom_dir)
        if not os.path.isdir(zoom_path):
            continue
        
        try:
            zoom = int(zoom_dir)
        except ValueError:
            continue
        
        zoom_tiles = 0
        for x_dir in os.listdir(zoom_path):
            x_path = os.path.join(zoom_path, x_dir)
            if not os.path.isdir(x_path):
                continue
            
            try:
                x = int(x_dir)
            except ValueError:
                continue
            
            for y_file in os.listdir(x_path):
                if not y_file.endswith('.png'):
                    continue
                
                try:
                    y = int(os.path.splitext(y_file)[0])
                    zoom_tiles += 1
                except ValueError:
                    continue
        
        tiles.append((zoom, zoom_tiles))
    
    return tiles

def retry_download_task(task_id, source_type="esri"):
    """使用新的瓦片源重试下载任务"""
    task_dir_name = f"task_{task_id}"
    task_dir = os.path.join(project_root, "data", "tiles", task_dir_name)
    
    if not os.path.exists(task_dir):
        logger.error(f"任务目录不存在: {task_dir}")
        return False
    
    # 创建下载器，使用新的瓦片源
    downloader = TileDownloader(
        output_dir=task_dir,
        max_workers=3,
        retries=5,
        timeout=15,
        source_type=source_type
    )
    
    # 先尝试查找任务的边界信息
    boundary_file = os.path.join(task_dir, "boundary.json")
    if os.path.exists(boundary_file):
        try:
            with open(boundary_file, 'r') as f:
                boundary = json.load(f)
            
            min_lat = boundary.get('min_lat')
            min_lon = boundary.get('min_lon')
            max_lat = boundary.get('max_lat')
            max_lon = boundary.get('max_lon')
            min_zoom = boundary.get('min_zoom')
            max_zoom = boundary.get('max_zoom')
            map_name = boundary.get('map_name')
            
            if None in (min_lat, min_lon, max_lat, max_lon, min_zoom, max_zoom):
                logger.warning("边界信息不完整，需要手动输入")
                return False
            
            logger.info(f"找到边界信息: {min_lat},{min_lon} - {max_lat},{max_lon}, 缩放级别: {min_zoom}-{max_zoom}")
            
            # 统一海域识别和推荐源
            tmp_downloader = TileDownloader()
            is_ocean = tmp_downloader.is_ocean_area(min_lat, min_lon, max_lat, max_lon)
            if is_ocean and source_type in ['osm', 'otm', 'carto']:
                logger.info("检测到海域任务，自动切换为esri_ocean优先")
                source_type = 'esri_ocean'
            
        except Exception as e:
            logger.error(f"读取边界信息失败: {str(e)}")
            return False
    else:
        logger.warning(f"未找到边界信息文件，无法自动重试")
        return False
    
    # 下载瓦片
    logger.info(f"使用 {source_type} 瓦片源重试下载任务 {task_id}")
    zoom_levels = range(min_zoom, max_zoom + 1)
    total_tiles, failed_tiles = downloader.download_area(
        min_lat=min_lat,
        min_lon=min_lon,
        max_lat=max_lat,
        max_lon=max_lon,
        zoom_levels=zoom_levels
    )
    
    # 如果下载成功，尝试转换为MBTiles文件
    if total_tiles > 0 and failed_tiles < total_tiles / 2:
        logger.info(f"下载完成，总瓦片数: {total_tiles}，失败瓦片数: {failed_tiles}")
        
        # 转换为MBTiles文件
        output_file = os.path.join(project_root, "data", "mbtiles", f"{map_name}.mbtiles")
        
        logger.info(f"将瓦片转换为MBTiles文件: {output_file}")
        converter = MBTilesConverter(
            tiles_dir=task_dir,
            output_file=output_file,
            name=map_name,
            description=f"{map_name} ({min_lat},{min_lon} - {max_lat},{max_lon})"
        )
        
        try:
            converted_tiles = converter.convert()
            logger.info(f"转换完成，共转换 {converted_tiles} 个瓦片")
            return True
        except Exception as e:
            logger.error(f"转换失败: {str(e)}")
            return False
    else:
        logger.warning(f"下载失败率过高: {failed_tiles}/{total_tiles}")
        return False

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='离线海图下载修复工具')
    parser.add_argument('--task', type=int, help='要修复的任务ID')
    parser.add_argument('--source', type=str, default='esri', choices=list(TILE_SOURCES.keys()), 
                        help='使用的瓦片源类型')
    parser.add_argument('--list', action='store_true', help='列出所有任务')
    parser.add_argument('--scan', type=int, help='扫描指定任务的瓦片文件')
    
    args = parser.parse_args()
    
    if args.list:
        # 列出所有任务
        task_dirs = list_task_dirs()
        if not task_dirs:
            print("未找到任何任务")
            return
        
        print(f"找到 {len(task_dirs)} 个任务:")
        for task_dir in task_dirs:
            task_id = task_dir.split('_')[1]
            tiles = scan_task_dir(task_dir)
            total_tiles = sum(t[1] for t in tiles)
            print(f"任务 #{task_id}: {total_tiles} 个瓦片")
            for zoom, count in tiles:
                print(f"  - 缩放级别 {zoom}: {count} 个瓦片")
            print("")
    
    elif args.scan is not None:
        # 扫描指定任务
        task_dir = f"task_{args.scan}"
        if not os.path.exists(os.path.join(project_root, "data", "tiles", task_dir)):
            print(f"错误: 任务 #{args.scan} 不存在")
            return
        
        tiles = scan_task_dir(task_dir)
        total_tiles = sum(t[1] for t in tiles)
        print(f"任务 #{args.scan} 包含 {total_tiles} 个瓦片:")
        for zoom, count in sorted(tiles):
            print(f"缩放级别 {zoom}: {count} 个瓦片")
    
    elif args.task is not None:
        # 修复指定任务
        success = retry_download_task(args.task, args.source)
        if success:
            print(f"任务 #{args.task} 修复成功")
        else:
            print(f"任务 #{args.task} 修复失败")
    
    else:
        parser.print_help()

if __name__ == "__main__":
    main() 